import torch
torch.set_num_threads(1)
import torch.nn as nn
from src.models.layer.embedding_layer import TabularEmbedding

class MLPRegressor(nn.Module):
    def __init__(self, num_cat_features, num_cont_features, num_classes=1,
                 embedding_dim=16, hidden_dim=128, hidden_dims=[256, 128], dropout=0.3):
        """
        :param num_cat_features: int，类别特征数量
        :param num_cont_features: int，连续特征数量
        :param num_classes: int，输出类别数量
        :param embedding_dim: int，类别embedding维度
        :param hidden_dim: int，连续特征投影维度(在其它模型表示隐藏层维度，这里做一个对齐)
        :param hidden_dims: list，隐藏层维度
        :param dropout: float，dropout率
        """
        super().__init__()
        
        self.embedding = TabularEmbedding(
            num_cat_features, num_cont_features, embedding_dim, hidden_dim, fuse=True
        )

        # 计算输入维度
        input_dim = num_cat_features * embedding_dim + hidden_dim

        layers = []
        for hidden_dim in hidden_dims:
            layers.append(nn.Linear(input_dim, hidden_dim))
            layers.append(nn.BatchNorm1d(hidden_dim))
            layers.append(nn.GELU())
            layers.append(nn.Dropout(dropout))
            input_dim = hidden_dim
        
        layers.append(nn.Linear(hidden_dims[-1], num_classes))  # 输出层

        self.regressor = nn.Sequential(*layers)

    def forward(self, batch):
        x_cat, x_cont, _ = batch
        x, _, _ = self.embedding(x_cat, x_cont)
        x = self.regressor(x)
        return x.squeeze(-1)  # [batch_size] 与 y 对齐
